// Copyright (C) joyinter
// All rights reserved.
// 
// Author : Jiang Li
// Date   : 2005-4-30

#include "stdafx.h"
#include "csumpeimage.h"
#include "crc.h"

namespace JoyGuard
{

//##ModelId=44CF17BF0007
CSumPeImage::CSumPeImage()
{
    m_error     = EC_OK;
    m_stopped   = false;
    m_process   = NULL;
    m_baseAddr  = NULL;
    m_readSize  = 64*1024;
    m_readDelay = 100;
    m_waitEvent = NULL;
    m_readed    = 0;
    m_ntHeaders = NULL;

    memset(&m_csumSect, 0, sizeof(m_csumSect));
}

//##ModelId=44CF17BF0002
void CSumPeImage::CheckStop()
{
    if(m_stopped)
    {
        throw ECheckStop();
    }
}

//##ModelId=44CF17BF0003
void CSumPeImage::ReadMemory(BYTE* dst, BYTE* src, int len)
{
    while(len > 0)
    {
        int c;
        if(len + m_readed >= m_readSize)
        {
            c = m_readSize-m_readed;
        }
        else
        {
            c = len;
        }
        SIZE_T readed = 0;
        if(!::ReadProcessMemory(m_process, src, dst, c, &readed))
        {
            m_error = EC_READ_MEMORY;
            throw ECheckStop();
        }
        m_readed += c;
        dst += c;
        src += c;
        len -= c;
        if(m_readed >= m_readSize)
        {
            m_readed -= m_readSize;
            ::WaitForSingleObject(m_waitEvent, m_readDelay);
            CheckStop();
        }
    }
}

//##ModelId=44CF17BF0008
bool CSumPeImage::Create(HANDLE waitEvent, HANDLE hProcess, BYTE* baseAddr, const CSumSect& ss, DWORD readSize, DWORD readDelay)
{
    CheckStop();

    m_process   = hProcess;
    m_baseAddr  = baseAddr;
    m_csumSect  = ss;
    m_waitEvent = waitEvent;
    if(readSize != 0)
    {
        m_readSize = readSize;
    }
    if(readDelay != 0)
    {
        m_readDelay = readDelay;
    }

    IMAGE_DOS_HEADER dosHeader;
    ReadMemory((BYTE*)&dosHeader, baseAddr, sizeof(dosHeader));

    IMAGE_NT_HEADERS ntHeaders;
    ReadMemory((BYTE*)&ntHeaders, baseAddr+dosHeader.e_lfanew, sizeof(ntHeaders));

    DWORD imgSize = ntHeaders.OptionalHeader.SizeOfImage;
    m_tempBase.resize(imgSize);
    BYTE* tempBase = &m_tempBase[0];
    ReadMemory(tempBase, baseAddr, imgSize);
    
    m_ntHeaders = (IMAGE_NT_HEADERS*)(tempBase+((IMAGE_DOS_HEADER*)tempBase)->e_lfanew);
    
    return true;
}

//##ModelId=44CF17BF0015
void CSumPeImage::Close()
{
    m_stopped = true;
}

//##ModelId=44CF17BF0016
bool CSumPeImage::Check()
{
    BYTE* tempBase = &m_tempBase[0];

    // read section data
    ReadMemory(tempBase+m_csumSect.m_rva, m_baseAddr+m_csumSect.m_rva, m_csumSect.m_size);

    // undo relocation
    DWORD relDelta = (DWORD)((UINT_PTR)tempBase - m_ntHeaders->OptionalHeader.ImageBase);
    IMAGE_BASE_RELOCATION* baseRel = (IMAGE_BASE_RELOCATION*)
        (tempBase + m_ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);
    while(baseRel->VirtualAddress)
    {
        BYTE* pPage  = (BYTE*)tempBase + baseRel->VirtualAddress;
        DWORD dwTypeCount = (baseRel->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION))/sizeof(WORD);
        baseRel = (IMAGE_BASE_RELOCATION*)((BYTE*)baseRel + baseRel->SizeOfBlock);

        if( baseRel->VirtualAddress < m_csumSect.m_rva ||
            baseRel->VirtualAddress > m_csumSect.m_rva + m_csumSect.m_size )
        {
            continue;
        }
        WORD* pType = (WORD*)(baseRel+1);
        WORD* pTypeMax  = pType + dwTypeCount;
        LONG Temp;
        for(; pType < pTypeMax; ++ pType)
        {
            PVOID pReloc = pPage + (*pType & 0xFFF);
            WORD fix = *pType >> 12;
            switch(fix)
            {
            case IMAGE_REL_BASED_ABSOLUTE:
                break;
            case IMAGE_REL_BASED_HIGH:
                *(WORD*)pReloc = *(WORD*)pReloc - HIWORD(relDelta);
                break;
            case IMAGE_REL_BASED_LOW:
                *(WORD*)pReloc = *(WORD*)pReloc - LOWORD(relDelta);
                break;
            case IMAGE_REL_BASED_HIGHLOW:
                *(DWORD*)pReloc -= relDelta;
                break;
            case IMAGE_REL_BASED_HIGHADJ:
                /*
                Temp = *(PUSHORT)pReloc << 16;
                ++ pType;
                Temp += (LONG)(*(PSHORT)pType);
                Temp += (ULONG) relDelta;
                Temp += 0x8000;
                *(PUSHORT)pReloc = (USHORT)(Temp >> 16);
                */
                break;
            case 11 : // IMAGE_REL_BASED_HIGH3ADJ:
               break;
            default:
                ::GetTickCount();
                break;
            }
        }// while(pType < pTypeMax)
    }// while(baseRel->VirtualAddress)

    CheckStop();

    const BYTE* sectBuf = tempBase+m_csumSect.m_rva;
    Crc csum;
    csum.Update(sectBuf, m_csumSect.m_size);
    if(m_csumSect.m_csum != csum.Final())
    {
        m_error = EC_CHECK_SUM;
        return false;
    }

    return true;
}

} // namespace JoyGuard
